#!/user/bin/env python3
# -*- coding: utf-8 -*-
import os,time,keyboard
os.system('title=ESP开发板工具&color 0a')  # 初始化控制台
os.system('cls')  # 清屏

# 获取当前时间
now = time.localtime(time.time())
# 格式化输出时间
ft=time.strftime("%Y_%m_%d", now)


list1=[]  # 存储找到的数据
list2=[]  # 存储错误信息

def connect_check(filename):  # 主板连接检查
    if os.path.exists(filename):
        os.remove('output.txt')
        print('请稍后，正在统计信息')
        os.system('esptool.exe flash_id >output.txt')
        os.system("attrib +H output.txt")
        os.system('cls')  # 清屏
    else:
        print('请稍后，正在统计信息')
        os.system('esptool.exe flash_id >output.txt')
        os.system("attrib +H output.txt")
        os.system('cls')  # 清屏
    file = open(filename, 'r')
    for line in reversed(list(open(filename))):
        if "Could not connect to" in line:  # 获取flash size，并存储在list中
            list2.append('Could not connect to')
    file.close()


def board_information(filename):  # 生成主板基础信息，后续提取相关信息并使用
    if os.path.exists(filename):
        os.remove('output.txt')
        print('请稍后，正在统计信息')
        os.system('esptool.exe flash_id >output.txt')
        os.system("attrib +H output.txt")
        os.system('cls')  # 清屏
    else:
        print('请稍后，正在统计信息')
        os.system('esptool.exe flash_id >output.txt')
        os.system("attrib +H output.txt")
        os.system('cls')  # 清屏
    file = open(filename, 'r')
    for line in reversed(list(open(filename))):
        if "Detected flash size" in line:  # 获取flash size，并存储在list中
            list1.append(line.rstrip().split(' ')[3])
    for line in reversed(list(open(filename))):
        if "MAC" in line:  # 获取MAC地址，并存储在list中
            # print(line.rstrip().split(' ')[1])
            list1.append(line.rstrip().split(' ')[1])
    for line in reversed(list(open(filename))):
        if "Crystal" in line:  # 获取晶振信息，并存储在list中
            list1.append(line.rstrip().split(' ')[2])
    for line in reversed(list(open(filename))):
        if "Features" in line:  # 获取特征信息，并存储在list中
            list1.append(line.rstrip().split(':')[1:][0][1:])
    for line in reversed(list(open(filename))):
        if "Chip" in line:  # 获取主板型号信息，并存储在list中
            list2 = line.rstrip().split(' ')[2:]
            for item in list2:
                list1.append(item)
    for line in reversed(list(open(filename))):
        if "Serial port" in line:  # 再找com口号，并存储在list中
            # print(line.rstrip().split(' ')[2])
            list1.append(line.rstrip().split(' ')[2])
    file.close()

def View_motherboard_information(filename):  # 查看主板信息
    board_information(filename)
    if len(list1) == 6:
        print('【主板相关信息】')
        print('内存大小：' + list1[0])
        print('MAC地址：' + list1[1])
        print('晶振：' + list1[2])
        print('主板特征：' + list1[3])
        print('主板型号：' + list1[4])
        print('使用的COM口：' + list1[5])
    elif len(list1) == 8:
        print('【主板相关信息】')
        print('内存大小：' + list1[0])
        print('MAC地址：' + list1[1])
        print('晶振：' + list1[2])
        print('主板特征：' + list1[3])
        print('主板型号：' + list1[4] + ' ' + list1[5] + ' ' + list1[6])
        print('使用的COM口：' + list1[7])
    else:
        pass
    print('')


def Backup_firmware(filename,name):  # 备份固件
    board_information(filename)
    if len(list1)==6:  # esp8266
        command1='esptool.exe --baud 115200 --port '+list1[5]+' read_flash 0x00000 0x400000 '+name+'_'+ft+'.bin'
        if list1[0]=='4MB':
            print(os.system(command1))
        else:
            pass
    elif len(list1)==8:  # esp32
        print('32')
        command1='esptool.exe --baud 115200 --port '+list1[7]+' read_flash 0x00000 0x400000 '+name+'_'+ft+'.bin'
        if list1[0]=='4MB':
            print(os.system(command1))
        else:
            pass

def recovery_firmware(filename,name):  # 擦除+恢复固件
    if os.path.exists(name+'.bin'):
        print('擦除+恢复模式即将自动开始，请确认原有相关数据已保存')
        Confirmation=input("原有相关数据已保存(y/n)：")
        if Confirmation=='y':
            os.system('cls')
            board_information(filename)
            Recover_filename=name+'.bin'
            if len(list1)==6:  # esp8266
                command1='esptool.exe --port '+list1[5]+' erase_flash'
                command2='esptool.exe --port '+list1[5]+' --baud 115200 write_flash --flash_size=detect 0 '+Recover_filename
                os.system(command1)
                os.system(command2)
            elif len(list1)==8:  # esp32
                command1='esptool.exe --port '+list1[7]+' erase_flash'
                command2='esptool.exe --port '+list1[7]+' --baud 115200 write_flash --flash_size=detect 0 '+Recover_filename
                os.system(command1)
                os.system(command2)
        elif Confirmation=='n':
            os.system('cls')
            print('n')
        else:
            os.system('cls')
            print('输入错误，请重新选择！')
            time.sleep(3)
            os.system('cls')
    else:
        print('备份文件未找到')
    if len(list1)==6:  # esp8266
        print(list1[5])
    elif len(list1)==8:  # esp32
        print(list1[7])


try:
    while 1:
        print('''***esp开发板工具***
请进行功能选择：
1、查看主板信息
2、固件备份
3、固件恢复
4、离开''')
        option = input("请选择（序号+回车）：")
        if option=='1':
            list2=[]
            connect_check("output.txt")
            if 'Could not connect to' in list2:
                print('设备未连接，请确认')
            else:
                os.system('cls')
                list1=[]
                View_motherboard_information("output.txt")
        elif option=='2':
            list2=[]
            connect_check("output.txt")
            if 'Could not connect to' in list2:
                print('设备未连接，请确认')
            else:
                os.system('cls')
                list1=[]
                name=input("请输入备份文件名：")
                Backup_firmware("output.txt",name)
        elif option=='3':
            list2=[]
            connect_check("output.txt")
            if 'Could not connect to' in list2:
                print('设备未连接，请确认')
            else:
                os.system('cls')
                list1=[]
                name=input("请输入固件文件名：")
                recovery_firmware("output.txt",name)
        elif option=='4':
            os.system('cls')
            if os.path.exists("output.txt"):
                os.remove('output.txt')
            break
        else:
            print('error')
except:
    os.system('cls')
    print('强制退出或异常中断，请联系管理员(QQ:1762103326)或站内私信！')
    os.remove('output.txt')
    print('')
    print('按任意键退出')
    os.system('pause > 1')
    os.system("attrib +H 1")
    os.remove('1')
    os.system('cls')